import unittest

from psychsim.pwl import *
from psychsim.reward import *
from psychsim.action import *
from psychsim.world import World,stateKey,actionKey,binaryKey,modelKey
from psychsim.agent import Agent

'''
This example examines partial order planning (putting socks and shoes on). Also considers case where Trevor believes his left shoe is already on.
'''

class ShoesTestCase(unittest.TestCase):
        
    def setUp(self):
        #basic setup
        self.world = World()
        self.trevor = Agent('trevor')
        self.world.addAgent(self.trevor)
        self.world.setOrder([self.trevor.name])

    def addStates(self):
        #define states
        self.world.defineState(self.trevor.name, 'left_foot', list, ['l_bare','l_sock','l_shoe'], description='left_foot state.')
        self.world.setState(self.trevor.name,'left_foot', 'l_bare')

        self.world.defineState(self.trevor.name, 'right_foot', list,['r_bare','r_sock','r_shoe'], description='right_foot state.')
        self.world.setState(self.trevor.name,'right_foot', 'r_bare')

    def addTermination(self):
        #set termination condition
        self.left = stateKey(self.trevor.name, 'left_foot')
        self.right = stateKey(self.trevor.name, 'right_foot')
        self.world.addTermination(makeTree({'if': equalRow(self.left, 'l_shoe'),
                                       True: {'if': equalRow(self.right, 'r_shoe'),
                                              True: True,
                                              False: False},
                                       False: False})) 

    def addActions(self):
        #add actions
        self.putl_sock = self.trevor.addAction({'verb': 'put on left sock'})
        self.putr_sock = self.trevor.addAction({'verb': 'put on right sock'})
        self.putl_shoe = self.trevor.addAction({'verb': 'put on left shoe'})
        self.putr_shoe = self.trevor.addAction({'verb': 'put on right shoe'})
        self.nop = self.trevor.addAction({'verb': 'do nothing'})
        
        #set legality
        self.trevor.setLegal(self.putl_sock, makeTree({'if': equalRow(self.left, 'l_bare'), True: True, False: False}))
        self.trevor.setLegal(self.putr_sock, makeTree({'if': equalRow(self.right, 'r_bare'), True: True, False: False}))
        self.trevor.setLegal(self.putl_shoe, makeTree({'if': equalRow(self.left, 'l_sock'), True: True, False: False}))
        self.trevor.setLegal(self.putr_shoe, makeTree({'if': equalRow(self.right, 'r_sock'), True: True, False: False}))
  
    def addDynamics(self):
        #define dynamics
        self.world.setDynamics(stateKey(self.trevor.name, 'left_foot'), self.putl_sock, makeTree(setToConstantMatrix(self.left, 'l_sock')))
        self.world.setDynamics(stateKey(self.trevor.name, 'right_foot'), self.putr_sock, makeTree(setToConstantMatrix(self.right, 'r_sock')))
        self.world.setDynamics(stateKey(self.trevor.name, 'left_foot'), self.putl_shoe, makeTree(setToConstantMatrix(self.left, 'l_shoe')))
        self.world.setDynamics(stateKey(self.trevor.name, 'right_foot'), self.putr_shoe, makeTree(setToConstantMatrix(self.right, 'r_shoe')))
        
    def addGoalsAndBeliefs(self):
        #goals
        goalRShoe = achieveFeatureValue(self.right, 'r_shoe')
        self.trevor.setReward(goalRShoe, 1.)

        self.trevor.setBelief(stateKey(self.trevor.name, 'left_foot'),'l_shoe', True)

    def setAttributes(self):
        #attributes
        self.trevor.setAttribute('rationality', 10.)        
        #self.trevor.setAttribute('selection', 'distribution')
        self.trevor.setAttribute('discount', 0.9)
        self.trevor.setAttribute('horizon', 6)

        self.world.setModel(self.trevor.name, True)
                           
    def testSetup(self):
        self.setUp()
        self.addStates()
        self.addTermination()
        self.addActions()
        self.addDynamics()
        self.addGoalsAndBeliefs()
        self.setAttributes()

    def testRun(self):
        self.testSetup()
        self.world.printState(beliefs=True)
        for round in range(6):
            self.world.explain(self.world.step(),level=3)
            self.world.printState(beliefs=True)
            #self.world.printModels()
            if self.world.terminated():
                break
    

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(ShoesTestCase):
       if method.startswith("test"):
          suite.addTest(ShoesTestCase(method))
    unittest.TextTestRunner().run(suite)